#include "wrapper_hypre.h"
#include <vector>

typedef HYPRE_Int idx_t;
typedef HYPRE_Real data_t;

const idx_t dims_per_blk[] = {
     64, 640, 768,
    //  64, 160, 192,
      1,   1,   1,
      1,   1,   1,
      1,   1,   1,
      1,   1,   1,
      1,   1,   1,
};

const idx_t offd_nnz [] = {
//  0   1   2   3   4   5
    0, 64, 64, 64, 64, 64,// 0
    64, 0,  0,  0,  0,  0,// 1
    64, 0,  0,  0,  0,  0,// 2
    64, 0,  0,  0,  0,  0,// 3
    64, 0,  0,  0,  0,  0,// 4
    64, 0,  0,  0,  0,  0,// 5
};

const idx_t num_diag = 7;

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);

    idx_t arg_cnt = 1;
    const idx_t pN[3] = { atoi(argv[arg_cnt++]), atoi(argv[arg_cnt++]),  atoi(argv[arg_cnt++]) };
    const std::string pathname(argv[arg_cnt++]);
    const std::string solver_name = std::string(argv[arg_cnt++]),
                        prec_name = std::string(argv[arg_cnt++]);
    const idx_t glb_nparts = sizeof(dims_per_blk) / sizeof(idx_t) / 3;
    idx_t blk_offsets[glb_nparts + 1];
    blk_offsets[0] = 0;
    for (idx_t b = 0; b < glb_nparts; b++) {
        blk_offsets[b + 1] = blk_offsets[b] 
            + dims_per_blk[b*3  ] * dims_per_blk[b*3+1] * dims_per_blk[b*3+2];
    }
    const idx_t tot_nrows = blk_offsets[glb_nparts];

    MPI_Comm blk0_comm = MPI_COMM_NULL;
    if (prec_name.substr(0, 5) == "SSAMG") {
        if (pN[0] * pN[1] * pN[2] != num_procs - 5) {
            if (my_pid == 0) printf("Error !! SSAMG requires 5 more procs\n");
            MPI_Abort(MPI_COMM_WORLD, -202);
        }
        MPI_Comm_split(MPI_COMM_WORLD, (my_pid < num_procs - 5) ? 1 : MPI_UNDEFINED, my_pid, & blk0_comm);
    } else {
        if (pN[0] * pN[1] * pN[2] != num_procs) {
            if (my_pid == 0) printf("Error !! Other MG requires no more procs\n");
            MPI_Abort(MPI_COMM_WORLD, -203);
        }
        blk0_comm = MPI_COMM_WORLD;
    }

    idx_t pCoords[3];
    pCoords[0] =  my_pid / (pN[1] * pN[2]);
    pCoords[1] = (my_pid - pCoords[0] * pN[1] * pN[2]) / pN[2];
    pCoords[2] =  my_pid - pCoords[0] * pN[1] * pN[2] - pCoords[1] * pN[2];

    idx_t begs[3], ends[3];
    for (idx_t d = 0; d < 3; d++) {
        idx_t loc_len = dims_per_blk[d] / pN[d];
        idx_t remain = dims_per_blk[d] - loc_len * pN[d];
        begs[d] = pCoords[d] * loc_len;
        if (pCoords[d] < remain) {
            begs[d] += pCoords[d];
            loc_len += 1;
        } else {
            begs[d] += remain;
        }
        ends[d] = begs[d] + loc_len;
    }

    std::vector<PartDescriptor<idx_t> > my_parts;
    if (prec_name.substr(0, 5) == "SSAMG") {
        const int np_blk0 = pN[0] * pN[1] * pN[2];
        if (my_pid == np_blk0) {
            PartDescriptor<idx_t> descr;
            descr.id = 1;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        } else if (my_pid == np_blk0 + 1) {
            PartDescriptor<idx_t> descr;
            descr.id = 2;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        } else if (my_pid == np_blk0 + 2) {
            PartDescriptor<idx_t> descr;
            descr.id = 3;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        } else if (my_pid == np_blk0 + 3) {
            PartDescriptor<idx_t> descr;
            descr.id = 4;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        } else if (my_pid == np_blk0 + 4) {
            PartDescriptor<idx_t> descr;
            descr.id = 5;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        } else {
            PartDescriptor<idx_t> descr;
            descr.id = 0;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) {
                descr.ilower[d] = begs[d];
                descr.iupper[d] = ends[d] - 1;
                descr.nelems *= (descr.iupper[d] - descr.ilower[d] + 1);
            }
            my_parts.emplace_back(descr);
        }
    }
    else {
        PartDescriptor<idx_t> descr;
        descr.id = 0;
        descr.nelems = 1;
        for (idx_t d = 0; d < 3; d++) {
            descr.ilower[d] = begs[d];
            descr.iupper[d] = ends[d] - 1;
            descr.nelems *= (descr.iupper[d] - descr.ilower[d] + 1);
        }
        my_parts.emplace_back(descr);
        if (begs[1] <= 319 && 319 < ends[1] &&
            begs[2] <= 383 && 383 < ends[2] && begs[0] == 0) {
        // if (begs[1] <= 79 && 79 < ends[1] &&
        //     begs[2] <= 95 && 95 < ends[2] && begs[0] == 0) {
            descr.id = 1;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        }
        if (begs[1] <=   0 &&   0 < ends[1] &&
            begs[2] <=   0 &&   0 < ends[2] && begs[0] == 0) {
            descr.id = 2;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        }
        // if (begs[1] <=  79 &&  79 < ends[1] &&
        if (begs[1] <= 319 && 319 < ends[1] &&
            begs[2] <=   0 &&   0 < ends[2] && begs[0] == 0) {
            descr.id = 3;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        }
        if (begs[1] <=   0 &&   0 < ends[1] &&
            // begs[2] <=  95 &&  95 < ends[2] && begs[0] == 0) {
            begs[2] <= 383 && 383 < ends[2] && begs[0] == 0) {
            descr.id = 4;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        }
        // if (begs[1] <= 159 && 159 < ends[1] &&
        //     begs[2] <= 191 && 191 < ends[2] && begs[0] == 0) {
        if (begs[1] <= 639 && 639 < ends[1] &&
            begs[2] <= 767 && 767 < ends[2] && begs[0] == 0) {
            descr.id = 5;
            descr.nelems = 1;
            for (idx_t d = 0; d < 3; d++) { descr.ilower[d] = descr.iupper[d] = 0; }
            my_parts.emplace_back(descr);
        }
    }
    const idx_t my_nparts = my_parts.size();
#ifdef DEBUG_PRINT
    for (int p = 0; p < num_procs; p++) {
        if (my_pid == p) {
            printf("Proc %d has %d parts\n", my_pid, my_nparts);
            for (idx_t i = 0; i < my_nparts; i++) {
                printf("  Part %d (%d,%d,%d) ~ (%d,%d,%d)\n", my_parts[i].id,
                    my_parts[i].ilower[0], my_parts[i].ilower[1], my_parts[i].ilower[2], 
                    my_parts[i].iupper[0], my_parts[i].iupper[1], my_parts[i].iupper[2] );
            }
            fflush(stdout);
        }
        MPI_Barrier(MPI_COMM_WORLD);
    }
#endif
    HYPRE_Init();
    idx_t obj_type = HYPRE_SSTRUCT;
    if (prec_name == "AMG") obj_type = HYPRE_PARCSR;
    if (my_pid == 0) {
        printf("object type : %d\n", obj_type);
    }

    HYPRE_SStructGrid ssgrid;
    HYPRE_SStructStencil stencil;
    HYPRE_SStructGraph ssgraph;
    HYPRE_SStructMatrix   A;
    HYPRE_SStructVector   b, x, y, x0;
    HYPRE_ParCSRMatrix par_A;
    HYPRE_ParVector par_b, par_x, par_y;
    TEST_RECORD records[TEST_CNT];
    TEST_CONFIG config;
    config.case_name = "oil";
    config.rtol = 1.0e-7;

    {
                // 创建网格
        HYPRE_SStructGridCreate(MPI_COMM_WORLD, 3, glb_nparts, &ssgrid);
        for (idx_t mb = 0; mb < my_nparts; mb++) {
            HYPRE_SStructGridSetExtents(ssgrid, my_parts[mb].id, my_parts[mb].ilower, my_parts[mb].iupper);
        }
        HYPRE_SStructVariable vtypes[1] = { HYPRE_SSTRUCT_VARIABLE_CELL };
        for (idx_t part = 0; part < glb_nparts; part++) {// 每个part都要执行定义变量的操作
            HYPRE_SStructGridSetVariables(ssgrid, part, sizeof(vtypes)/sizeof(HYPRE_SStructVariable), vtypes);
        }
        HYPRE_SStructGridAssemble(ssgrid);// a collective call finalizing the grid assembly.
        // 创建stencil
        HYPRE_SStructStencilCreate(3, num_diag, &stencil);
        idx_t offsets[7][3] = {
            {-1,0,0}, {0,-1,0}, {0,0,-1}, {0,0,0}, {0,0,1}, {0,1,0}, {1,0,0},// 正确
            // {-1,0,0}, {0,0,-1}, {0,-1,0}, {0,0,0}, {0,1,0}, {0,0,1}, {1,0,0},// 差一些
            // {0,0,-1}, {-1,0,0}, {0,-1,0}, {0,0,0}, {0,1,0}, {1,0,0}, {0,0,1},// 差很多
            // {0,0,-1}, {0,-1,0}, {-1,0,0}, {0,0,0}, {1,0,0}, {0,1,0}, {0,0,1},// 差很多
            // {0,-1,0}, {0,0,-1}, {-1,0,0}, {0,0,0}, {1,0,0}, {0,0,1}, {0,1,0},// 差很多
            // {0,-1,0}, {-1,0,0}, {0,0,-1}, {0,0,0}, {0,0,1}, {1,0,0}, {0,1,0},// 差很多

            // {0,0,1}, {0,1,0}, {1,0,0}, {0,0,0}, {-1,0,0}, {0,-1,0}, {0,0,-1}, 
        };
        for (idx_t e = 0; e < num_diag; e++) {
            HYPRE_SStructStencilSetEntry(stencil, e, offsets[e], 0);// 只有0号变量一种
        }
        // 创建图
        HYPRE_SStructGraphCreate(MPI_COMM_WORLD, ssgrid, &ssgraph);
        HYPRE_SStructGraphSetObjectType(ssgraph, obj_type);
        for (idx_t part = 0; part < glb_nparts; part++) {// 每个part都要执行定义模板的操作
            HYPRE_SStructGraphSetStencil(ssgraph, part, 0, stencil);
        }
        // 添加上非stencil的连接关系
        for (idx_t mp = 0; mp < my_nparts; mp++) {
            const PartDescriptor<idx_t> & loc_part = my_parts[mp];
            // TODO: 在目录下寻找名为"B%d_Offd*"的文件，并从文件大小确定有多少个非零元
            FILE * fp = nullptr;
            for (idx_t np = 0; np < glb_nparts; np++) {// 遍历每一个可能的邻居块
                const idx_t nnz = offd_nnz[loc_part.id * glb_nparts + np];
                if (nnz == 0) continue;
                // 读入该非对角块的数据

                std::vector<COO<idx_t, data_t> > buf(nnz);
                fp = fopen((pathname + "/B"+std::to_string(loc_part.id)+"_Offd"+std::to_string(np)).c_str(), "rb");
                size_t ret = fread(buf.data(), sizeof(COO<idx_t, data_t>), nnz, fp);
                assert(ret == (size_t) nnz);
                fclose(fp);

                idx_t loc_dims[3], ngb_dims[3];
                for (idx_t d = 0; d < 3; d++) {
                    loc_dims[d] = dims_per_blk[loc_part.id * 3 + d];
                    ngb_dims[d] = dims_per_blk[     np     * 3 + d];
                }
                const idx_t loc_nelems = loc_dims[0] * loc_dims[1] * loc_dims[2],
                            ngb_nelems = ngb_dims[0] * ngb_dims[1] * ngb_dims[2];
                // 为该非对角块设置连接关系
                for (idx_t k = 0; k < nnz; k++) {
                    const idx_t loc_1D = buf[k].r - blk_offsets[loc_part.id],// 该非零元的行序号在本块内的一维索引
                                ngb_1D = buf[k].c - blk_offsets[     np    ];// 该非零元的列序号在邻居块内的一维索引
                    assert(loc_1D >= 0 && loc_1D < loc_nelems);
                    assert(ngb_1D >= 0 && ngb_1D < ngb_nelems);
                    idx_t index[3], to_id[3];
                    index[0] = (loc_1D                                       ) / (loc_dims[1] * loc_dims[2]);
                    index[1] = (loc_1D - index[0] * loc_dims[1] * loc_dims[2]) / (loc_dims[2]              );
                    index[2] = (loc_1D - index[0] * loc_dims[1] * loc_dims[2] - index[1] * loc_dims[2]);
                    to_id[0] = (ngb_1D                                       ) / (ngb_dims[1] * ngb_dims[2]);
                    to_id[1] = (ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2]) / (ngb_dims[2]              );
                    to_id[2] = (ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2] - to_id[1] * ngb_dims[2]);

                    if (loc_part.id == 0) {
                        assert(ngb_dims[0] == 1 && ngb_dims[1] == 1 && ngb_dims[2] == 1);
                        if (loc_part.ilower[0] <= index[0] && index[0] <= loc_part.iupper[0] &&
                            loc_part.ilower[1] <= index[1] && index[1] <= loc_part.iupper[1] &&
                            loc_part.ilower[2] <= index[2] && index[2] <= loc_part.iupper[2] )// 确实属于自己的负责范围
                        {
                            // printf("Add ent %d (%d,%d,%d) -- %d (%d,%d,%d)\n", 
                            //     loc_part.id, index[0], index[1], index[2], 
                            //     np, to_id[0], to_id[1], to_id[2]);
                            HYPRE_SStructGraphAddEntries(ssgraph, loc_part.id, index, 0, np, to_id, 0);
                        }
                    }
                    else {
                        assert(loc_dims[0] == 1 && loc_dims[1] == 1 && loc_dims[2] == 1);
                        // printf("Add ent %d (%d,%d,%d) -- %d (%d,%d,%d)\n", 
                        //         loc_part.id, index[0], index[1], index[2], 
                        //         np, to_id[0], to_id[1], to_id[2]);
                        HYPRE_SStructGraphAddEntries(ssgraph, loc_part.id, index, 0, np, to_id, 0);
                    }
                }// nnz loop
            }// ngb part loop
        }// my part loop
        HYPRE_SStructGraphAssemble(ssgraph);

        MPI_Barrier(MPI_COMM_WORLD);
        if (my_pid == 0) printf("Done assem graph\n");

        // 建立向量
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &b);// Create an empty vector object
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &x);
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &y);
        HYPRE_SStructVectorCreate(MPI_COMM_WORLD, ssgrid, &x0);
        HYPRE_SStructVectorSetObjectType(b, obj_type);// Set the object type for the vectors to be the same as was already set for the matrix
        HYPRE_SStructVectorSetObjectType(x, obj_type);
        HYPRE_SStructVectorSetObjectType(y, obj_type);
        HYPRE_SStructVectorSetObjectType(x0, obj_type);
        HYPRE_SStructVectorInitialize(b);// Indicate that the vector coefficients are ready to be set
        HYPRE_SStructVectorInitialize(x);
        HYPRE_SStructVectorInitialize(y);
        HYPRE_SStructVectorInitialize(x0);

        data_t my_x_dot = 0.0, my_b_dot = 0.0;

        {// 填充向量数据
            FILE * fp = nullptr; size_t ret;
            for (idx_t mp = 0; mp < my_nparts; mp++) {
                PartDescriptor<idx_t> & loc_part = my_parts[mp];
                idx_t dims[3];
                for (idx_t d = 0; d < 3; d++)
                    dims[d] = loc_part.iupper[d] - loc_part.ilower[d] + 1;
                data_t * vec_buf = new data_t [loc_part.nelems];// [dim0][dim1][dim2]
                data_t * set_buf = new data_t [loc_part.nelems];// [dim2][dim1][dim0]

                if (loc_part.id == 0) {
                    MPI_Datatype mpi_type = MPI_DATATYPE_NULL;
                    idx_t sizes[3], subsizes[3], starts[3];
                    MPI_Status status;
                    for (idx_t d = 0; d < 3; d++) {
                        sizes   [d] = dims_per_blk[loc_part.id * 3 + d];
                        subsizes[d] = loc_part.iupper[d] - loc_part.ilower[d] + 1;
                        starts  [d] = loc_part.ilower[d];
                    }
                    int tot_len = subsizes[0] * subsizes[1] * subsizes[2];
                    static_assert(sizeof(data_t) == 8);
                    MPI_Type_create_subarray(3, sizes, subsizes, starts, MPI_ORDER_C, MPI_DOUBLE, &mpi_type);
                    MPI_Type_commit(&mpi_type);
                    MPI_File fh = MPI_FILE_NULL;
                    int ret = MPI_File_open(blk0_comm, (pathname + "/b.B0").c_str(), MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
                    if (ret != MPI_SUCCESS) printf("Could not open file: ret %d\n", ret);
                    
                    MPI_File_set_view(fh, 0, MPI_DOUBLE, mpi_type, "native", MPI_INFO_NULL);
                    MPI_File_read_all(fh, vec_buf, tot_len, MPI_DOUBLE, &status);
                    ret = MPI_File_close(&fh);
                    if (ret != MPI_SUCCESS) printf("Could not close file: ret %d\n", ret);
                    
                    #pragma omp parallel for collapse(2) schedule(static) reduction(+:my_b_dot)
                    for (idx_t i0 = 0; i0 < dims[0]; i0++)
                    for (idx_t i1 = 0; i1 < dims[1]; i1++)
                    for (idx_t i2 = 0; i2 < dims[2]; i2++) {
                        data_t tmp = vec_buf[((i0 * dims[1]) + i1) * dims[2] + i2];
                        set_buf[((i2 * dims[1]) + i1) * dims[0] + i0] = tmp;
                        my_b_dot += tmp * tmp;
                        // if (dims[0] == 1) printf("b %d %.15e\n", loc_part.id, tmp);
                    }
                    HYPRE_SStructVectorSetBoxValues(b, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);

                    ret = MPI_File_open(blk0_comm, (pathname + "/x0.B0").c_str(), MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
                    if (ret != MPI_SUCCESS) printf("Could not open file: ret %d\n", ret);
                    MPI_File_set_view(fh, 0, MPI_DOUBLE, mpi_type, "native", MPI_INFO_NULL);
                    MPI_File_read_all(fh, vec_buf, tot_len, MPI_DOUBLE, &status);
                    ret = MPI_File_close(&fh);
                    if (ret != MPI_SUCCESS) printf("Could not close file: ret %d\n", ret);

                    #pragma omp parallel for collapse(2) schedule(static) reduction(+:my_x_dot)
                    for (idx_t i0 = 0; i0 < dims[0]; i0++)
                    for (idx_t i1 = 0; i1 < dims[1]; i1++)
                    for (idx_t i2 = 0; i2 < dims[2]; i2++) {
                        data_t tmp = vec_buf[((i0 * dims[1]) + i1) * dims[2] + i2];
                        set_buf[((i2 * dims[1]) + i1) * dims[0] + i0] = tmp;
                        my_x_dot += tmp * tmp;
                        // if (dims[0] == 1) printf("x %d %.15e\n", loc_part.id, tmp);
                    }
                    HYPRE_SStructVectorSetBoxValues(x, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);

                    MPI_Type_free(&mpi_type);
                }
                else {
                    assert(dims[0] == 1 && dims[1] == 1 && dims[2] == 1);
                    fp = fopen((pathname + "/b.B" + std::to_string(loc_part.id)).c_str(), "rb");
                    ret = fread(vec_buf, sizeof(data_t), loc_part.nelems, fp); assert(ret == (size_t) loc_part.nelems);
                    fclose(fp);
                    #pragma omp parallel for collapse(2) schedule(static)
                    for (idx_t i0 = 0; i0 < dims[0]; i0++)
                    for (idx_t i1 = 0; i1 < dims[1]; i1++)
                    for (idx_t i2 = 0; i2 < dims[2]; i2++) {
                        data_t tmp = vec_buf[((i0 * dims[1]) + i1) * dims[2] + i2];
                        set_buf[((i2 * dims[1]) + i1) * dims[0] + i0] = tmp;
                        my_b_dot += tmp * tmp;
                        // if (dims[0] == 1) printf("b %d %.15e\n", loc_part.id, tmp);
                    }
                    HYPRE_SStructVectorSetBoxValues(b, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);

                    fp = fopen((pathname + "/x0.B" + std::to_string(loc_part.id)).c_str(), "rb");
                    ret = fread(vec_buf, sizeof(data_t), loc_part.nelems, fp); assert(ret == (size_t) loc_part.nelems);
                    fclose(fp);
                    #pragma omp parallel for collapse(2) schedule(static)
                    for (idx_t i0 = 0; i0 < dims[0]; i0++)
                    for (idx_t i1 = 0; i1 < dims[1]; i1++)
                    for (idx_t i2 = 0; i2 < dims[2]; i2++) {
                        data_t tmp = vec_buf[((i0 * dims[1]) + i1) * dims[2] + i2];
                        set_buf[((i2 * dims[1]) + i1) * dims[0] + i0] = tmp;
                        my_x_dot += tmp * tmp;
                        // if (dims[0] == 1) printf("x %d %.15e\n", loc_part.id, tmp);
                    }
                    HYPRE_SStructVectorSetBoxValues(x, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);
                }
                
                #pragma omp parallel for schedule(static)
                for (idx_t i = 0; i < loc_part.nelems; i++) set_buf[i] = 0.0;
                HYPRE_SStructVectorSetBoxValues(y, loc_part.id, loc_part.ilower, loc_part.iupper, 0, set_buf);

                delete [] vec_buf; vec_buf = nullptr;
                delete [] set_buf; set_buf = nullptr;
            }// my part loop
        }// fill in vector data
        HYPRE_SStructVectorAssemble(b);
        HYPRE_SStructVectorAssemble(x);
        HYPRE_SStructVectorAssemble(y);// 对y也要创建
        HYPRE_SStructVectorAssemble(x0);
        HYPRE_SStructVectorCopy(x, x0);

        MPI_Barrier(MPI_COMM_WORLD);
        if (my_pid == 0) printf("Done fill vec\n");
        
        HYPRE_SStructMatrixCreate(MPI_COMM_WORLD, ssgraph, &A);// Create an empty matrix object
        HYPRE_SStructMatrixSetObjectType(A, obj_type);
        HYPRE_SStructMatrixInitialize(A);// Indicate that the matrix coefficients are ready to be set 
        idx_t stencil_indices[num_diag];
        for (idx_t j = 0; j < num_diag; j++)
            stencil_indices[j] = j;
        {// 填充矩阵数据
            FILE * fp = nullptr; size_t ret;
            for (idx_t mp = 0; mp < my_nparts; mp++) {
                PartDescriptor<idx_t> & loc_part = my_parts[mp];
                idx_t dims[3];
                for (idx_t d = 0; d < 3; d++)
                    dims[d] = loc_part.iupper[d] - loc_part.ilower[d] + 1;
                data_t * mat_buf = new data_t [loc_part.nelems * num_diag];
                data_t * set_buf = new data_t [loc_part.nelems * num_diag];
                if (loc_part.id == 0) {
                    // 先设置对角块的数值
                    MPI_Datatype mpi_type = MPI_DATATYPE_NULL;
                    idx_t sizes[4], subsizes[4], starts[4];
                    MPI_Status status;
                    for (idx_t d = 0; d < 3; d++) {
                        sizes   [d] = dims_per_blk[loc_part.id * 3 + d];
                        subsizes[d] = loc_part.iupper[d] - loc_part.ilower[d] + 1;
                        starts  [d] = loc_part.ilower[d];
                    }
                    sizes[3] = subsizes[3] = num_diag;
                    starts[3] = 0;
                    int tot_len = subsizes[0] * subsizes[1] * subsizes[2] * subsizes[3];
                    static_assert(sizeof(data_t) == 8);
                    MPI_Type_create_subarray(4, sizes, subsizes, starts, MPI_ORDER_C, MPI_DOUBLE, &mpi_type);
                    MPI_Type_commit(&mpi_type);

                    MPI_File fh = MPI_FILE_NULL;
                    int ret = MPI_File_open(blk0_comm, (pathname + "/B0_Diag").c_str(), MPI_MODE_RDONLY, MPI_INFO_NULL, &fh);
                    if (ret != MPI_SUCCESS) printf("Could not open file: ret %d\n", ret);
                    
                    MPI_File_set_view(fh, 0, MPI_DOUBLE, mpi_type, "native", MPI_INFO_NULL);
                    MPI_File_read_all(fh, mat_buf, tot_len, MPI_DOUBLE, &status);
                    ret = MPI_File_close(&fh);
                    if (ret != MPI_SUCCESS) printf("Could not close file: ret %d\n", ret);

                    #pragma omp parallel for collapse(2) schedule(static)
                    for (idx_t i0 = 0; i0 < dims[0]; i0++)
                    for (idx_t i1 = 0; i1 < dims[1]; i1++)
                    for (idx_t i2 = 0; i2 < dims[2]; i2++) {
                        data_t       * dst_ptr = set_buf + (((i2 * dims[1]) + i1) * dims[0] + i0) * num_diag;
                        const data_t * src_ptr = mat_buf + (((i0 * dims[1]) + i1) * dims[2] + i2) * num_diag;
                        for (idx_t d = 0; d < num_diag; d++) {
                            dst_ptr[d] = src_ptr[d];
                            assert(dst_ptr[d] == dst_ptr[d]);
                        }
                    }
                    HYPRE_SStructMatrixSetBoxValues(A, loc_part.id, loc_part.ilower, loc_part.iupper, 0,
                        num_diag, stencil_indices, set_buf);

                    MPI_Type_free(&mpi_type);
                }
                else {
                    assert(dims[0] == 1 && dims[1] == 1 && dims[2] == 1);
                    fp = fopen((pathname + "/B" + std::to_string(loc_part.id) + "_Diag").c_str(), "rb");
                    ret = fread(mat_buf, sizeof(data_t), loc_part.nelems * num_diag, fp); assert(ret == (size_t)(loc_part.nelems * num_diag));
                    fclose(fp);
                    #pragma omp parallel for collapse(2) schedule(static)
                    for (idx_t i0 = 0; i0 < dims[0]; i0++)
                    for (idx_t i1 = 0; i1 < dims[1]; i1++)
                    for (idx_t i2 = 0; i2 < dims[2]; i2++) {
                        data_t       * dst_ptr = set_buf + (((i2 * dims[1]) + i1) * dims[0] + i0) * num_diag;
                        const data_t * src_ptr = mat_buf + (((i0 * dims[1]) + i1) * dims[2] + i2) * num_diag;
                        for (idx_t d = 0; d < num_diag; d++) {
                            dst_ptr[d] = src_ptr[d];
                            // if (dst_ptr[d] != dst_ptr[d]) {
                            //     printf("Error (%d,%d,%d): d %d val %.3e \n", i0, i1, i2, d, dst_ptr[d]);
                            // }
                            assert(dst_ptr[d] == dst_ptr[d]);
                        }
                    }
                    HYPRE_SStructMatrixSetBoxValues(A, loc_part.id, loc_part.ilower, loc_part.iupper, 0,
                        num_diag, stencil_indices, set_buf);
                }
                delete [] mat_buf; mat_buf = nullptr;
                delete [] set_buf; set_buf = nullptr;
                
                // 再逐个设置非对角块的数值
                idx_t loc_dims[3];
                for (idx_t d = 0; d < 3; d++) loc_dims[d] = dims_per_blk[loc_part.id * 3 + d];
                idx_t * nentries = new idx_t [loc_part.nelems];// 记录已经为该位置设置了多少个非零元
                #pragma omp parallel for schedule(static)
                for (idx_t i = 0; i < loc_part.nelems; i++) nentries[i] = num_diag;

                for (idx_t ngb = 0; ngb < glb_nparts; ngb++) {// 遍历每一个可能的邻居块
                    const idx_t nnz = offd_nnz[loc_part.id * glb_nparts + ngb];
                    if (nnz == 0) continue;
                    // 读入该非对角块的数据
                    std::vector<COO<idx_t, data_t> > buf(nnz);
                    fp = fopen((pathname + "/B"+std::to_string(loc_part.id)+"_Offd"+std::to_string(ngb)).c_str(), "rb");
                    size_t ret = fread(buf.data(), sizeof(COO<idx_t, data_t>), nnz, fp);
                    assert(ret == (size_t) nnz);
                    fclose(fp);
                    // 为该非对角块设置连接关系
                    
                    for (idx_t k = 0; k < nnz; k++) {
                        const idx_t blk_1D = buf[k].r - blk_offsets[loc_part.id];// 该非零元的行序号在本块内的一维索引
                        idx_t index[3];// 块内坐标
                        index[0] = (blk_1D                                       ) / (loc_dims[1] * loc_dims[2]);
                        index[1] = (blk_1D - index[0] * loc_dims[1] * loc_dims[2]) / (loc_dims[2]              );
                        index[2] = (blk_1D - index[0] * loc_dims[1] * loc_dims[2] - index[1] * loc_dims[2]);

                        if (loc_part.ilower[0] <= index[0] && index[0] <= loc_part.iupper[0] &&
                            loc_part.ilower[1] <= index[1] && index[1] <= loc_part.iupper[1] &&
                            loc_part.ilower[2] <= index[2] && index[2] <= loc_part.iupper[2] )// 确实属于自己的负责范围
                        {
                            const idx_t loc_1D = ((index[0] - loc_part.ilower[0]) * dims[1] 
                                                + (index[1] - loc_part.ilower[1]))* dims[2]
                                                +  index[2] - loc_part.ilower[2];
                            // printf("Set val %d (%d,%d,%d) -- %.3e  %d\n", 
                            //     loc_part.id, index[0], index[1], index[2], buf[k].v, nentries[loc_1D]);
                            assert(buf[k].v == buf[k].v);
                            HYPRE_SStructMatrixSetValues(A, loc_part.id, index, 0, 1, &nentries[loc_1D], &buf[k].v);
                            nentries[loc_1D] ++;// 计数增加
                        }
                    }// nnz loop
                }// ngb loop
                delete [] nentries; nentries = nullptr;
            }// my part loop
        }// fill in matrix data
        HYPRE_SStructMatrixAssemble(A);// a collective call finalizing the matrix assembly.
        
        for (idx_t test = 0; test < TEST_CNT; test++) {
            HYPRE_SStructVectorCopy(x0, x);

            data_t b_dot, x_dot, Ab_dot, Ax_dot;
            idx_t ret;
            // 做spmv检验一下数据传对没有
            if (obj_type == HYPRE_PARCSR) {
                HYPRE_SStructMatrixGetObject(A, (void **) &par_A);
                HYPRE_SStructVectorGetObject(b, (void **) &par_b);
                HYPRE_SStructVectorGetObject(x, (void **) &par_x);
                HYPRE_SStructVectorGetObject(y, (void **) &par_y);
                ret = HYPRE_ParVectorInnerProd(par_b, par_b, &b_dot);
                ret = HYPRE_ParVectorInnerProd(par_x, par_x, &x_dot);
                HYPRE_ParCSRMatrixMatvec(1.0, par_A, par_b, 0.0, par_y);
                ret = HYPRE_ParVectorInnerProd(par_y, par_y, &Ab_dot);
                HYPRE_ParCSRMatrixMatvec(1.0, par_A, par_x, 0.0, par_y);
                ret = HYPRE_ParVectorInnerProd(par_y, par_y, &Ax_dot);
            } else { assert(obj_type == HYPRE_SSTRUCT);
                ret = HYPRE_SStructInnerProd(b, b, &b_dot);
                ret = HYPRE_SStructInnerProd(x, x, &x_dot);
                HYPRE_SStructMatrixMatvec(1.0, A, b, 0.0, y);
                ret = HYPRE_SStructInnerProd(y, y, &Ab_dot);
                HYPRE_SStructMatrixMatvec(1.0, A, x, 0.0, y);
                ret = HYPRE_SStructInnerProd(y, y, &Ax_dot);
            }

            if (my_pid == 0) {
                // printf("(myb, myb) = %.20e\n",  (double)my_b_dot);
                // printf("(myx, myx) = %.20e\n",  (double)my_x_dot);
                printf("(  b,   b) = %.20e\n",  (double)b_dot);
                printf("(  x,   x) = %.20e\n",  (double)x_dot);
                printf("(A*b, A*b) = %.20e\n", (double)Ab_dot);
                printf("(A*x, A*x) = %.20e\n", (double)Ax_dot);
            }

            
            // 计算初始残差以便确定收敛阈值
            data_t r_nrm2 = 0.0, b_nrm2 = 0.0;
            // check_residual(A, x, b, y, r_nrm2, b_nrm2);
            // config.atol = r_nrm2 * 1.0e-6;

            buildup_solver(solver_name, prec_name, config);

            idx_t num_iterations = 0;
            data_t t_setup = 0.0, t_solve = 0.0, final_res_norm = 0.0;

            if (obj_type == HYPRE_SSTRUCT)
                setup_and_solve(solver_name, prec_name, A, b, x, t_setup, t_solve, final_res_norm, num_iterations);
            else if (obj_type == HYPRE_PARCSR) {
                setup_and_solve(solver_name, prec_name, par_A, par_b, par_x, t_setup, t_solve, final_res_norm, num_iterations);
                HYPRE_SStructVectorGather(b);
                HYPRE_SStructVectorGather(x);
            }

            // 计算真实残差
            check_residual(prec_name, A, x, b, y, r_nrm2, b_nrm2);
            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", r_nrm2, r_nrm2/b_nrm2);
                printf("Iterations = %d\n", num_iterations);
                printf("Time cost %.5f %.5f %.5f %d\n", (double)t_setup, (double)t_solve, (double)(t_setup + t_solve), num_iterations);
                printf("Final Relative Residual Norm = %e\n", (double)final_res_norm);
                printf("\n");
            }

            records[test].iter = num_iterations;
            records[test].setup = t_setup;
            records[test].solve = t_solve;

            destroy_solver(solver_name, prec_name);
        }// test loop

        if (my_pid == 0) {
            std::sort(records, records + TEST_CNT);
            const TEST_RECORD & best = records[0], & worst = records[TEST_CNT - 1];
            TEST_RECORD avg;
            for (idx_t i = 0; i < TEST_CNT; i++) {
                avg.setup += records[i].setup;
                avg.solve += records[i].solve;
                avg.prec  += records[i].prec;
                avg.iter  += records[i].iter;
            }
            avg.setup /= TEST_CNT;
            avg.solve /= TEST_CNT;
            avg.prec  /= TEST_CNT;
            avg.iter  /= TEST_CNT;
            data_t min_tot = best .setup + best .solve;
            data_t max_tot = worst.setup + worst.solve;
            data_t avg_tot = avg  .setup + avg  .solve;
            printf("\n  Summary of %d tests:\n", TEST_CNT);
            printf("     Setup time  Solve time  #Iter  Total time\n");
            printf("Min    %.6f    %.6f  %5d  %.6f\n", best .setup, best .solve, best .iter, min_tot);
            printf("Avg    %.6f    %.6f  %5d  %.6f\n", avg  .setup, avg  .solve, avg  .iter, avg_tot);
            printf("Max    %.6f    %.6f  %5d  %.6f\n", worst.setup, worst.solve, worst.iter, max_tot);    
        }

        HYPRE_SStructMatrixDestroy(A);
        HYPRE_SStructVectorDestroy(b); HYPRE_SStructVectorDestroy(x); HYPRE_SStructVectorDestroy(y); HYPRE_SStructVectorDestroy(x0); 
        HYPRE_SStructGraphDestroy(ssgraph);
        HYPRE_SStructStencilDestroy(stencil);
        HYPRE_SStructGridDestroy(ssgrid);
    }
    HYPRE_Finalize();
    MPI_Finalize();
    return 0;
}